import mobi
import os
import re
from bs4 import BeautifulSoup
import html2text
import shutil
from PIL import Image
import io

def clean_text(text):
    """清理文本，移除多余的空格和换行"""
    # 移除多余的空格
    text = re.sub(r'\s+', ' ', text)
    # 移除多余的空行
    text = re.sub(r'\n\s*\n', '\n\n', text)
    return text.strip()

def extract_chapters(html_content):
    """从 HTML 内容中提取章节"""
    soup = BeautifulSoup(html_content, 'html.parser')
    chapters = []
    
    # 查找所有可能的章节标题标签
    for tag in soup.find_all(['h1', 'h2', 'h3']):
        chapter = {'title': tag.get_text().strip(), 'content': []}
        current = tag.next_sibling
        
        # 收集直到下一个标题的内容
        while current and current.name not in ['h1', 'h2', 'h3']:
            if current.name:
                chapter['content'].append(str(current))
            current = current.next_sibling
        
        chapters.append(chapter)
    
    return chapters

def mobi_to_markdown(mobi_path, output_dir=None):
    """
    将 MOBI 文件转换为 Markdown 格式，每个章节保存为单独的文件
    
    Args:
        mobi_path: MOBI 文件路径
        output_dir: 输出目录路径，如果为 None 则使用 MOBI 文件名创建目录
    """
    if not os.path.exists(mobi_path):
        print(f"错误：文件 {mobi_path} 不存在")
        return
    
    # 创建输出目录
    if output_dir is None:
        output_dir = os.path.splitext(mobi_path)[0]
    os.makedirs(output_dir, exist_ok=True)
    
    try:
        # 解压 MOBI 文件
        tempdir, filepath = mobi.extract(mobi_path)
        
        # 读取 HTML 内容
        with open(filepath, 'r', encoding='utf-8') as f:
            html_content = f.read()
        
        # 创建转换器
        h = html2text.HTML2Text()
        h.ignore_links = False
        h.ignore_images = False
        h.body_width = 0  # 不限制行宽
        
        # 提取章节
        chapters = extract_chapters(html_content)
        
        if not chapters:
            print("警告：未找到章节，将整个内容保存为一个文件")
            # 如果没有找到章节，将整个内容保存为一个文件
            markdown_content = h.handle(html_content)
            markdown_content = clean_text(markdown_content)
            
            output_path = os.path.join(output_dir, "content.md")
            with open(output_path, 'w', encoding='utf-8') as f:
                f.write(markdown_content)
        else:
            # 保存每个章节
            for i, chapter in enumerate(chapters, 1):
                # 创建安全的文件名
                safe_filename = re.sub(r'[\\/*?:"<>|]', '_', chapter['title'])
                safe_filename = re.sub(r'\s+', '_', safe_filename)
                
                # 转换章节内容为 Markdown
                chapter_content = '\n'.join(chapter['content'])
                markdown_content = h.handle(chapter_content)
                markdown_content = clean_text(markdown_content)
                
                # 保存章节内容
                output_path = os.path.join(output_dir, f"{i:02d}_{safe_filename}.md")
                with open(output_path, 'w', encoding='utf-8') as f:
                    f.write(f"# {chapter['title']}\n\n")
                    f.write(markdown_content)
                
                print(f"已保存章节：{chapter['title']}")
        
        # 清理临时文件
        shutil.rmtree(tempdir)
        
        print(f"\n转换完成！所有章节已保存至目录：{output_dir}")
    
    except Exception as e:
        print(f"转换过程中出现错误：{str(e)}")

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='将 MOBI 文件转换为 Markdown 格式，每个章节保存为单独的文件')
    parser.add_argument('mobi_path', help='MOBI 文件路径')
    parser.add_argument('-o', '--output', help='输出目录路径（可选）')
    
    args = parser.parse_args()
    mobi_to_markdown(args.mobi_path, args.output) 